from win32com.client import Dispatch
import networkx as nx
import pickle
catGeoTypePoint2D_ = 2
catGeoTypeLine2D_ = 3


CATIA = Dispatch("CATIA.Application")
vba_function_name = 'get_coordinates'
vba_code = """
Public Function get_coordinates(point_2d_instance)
    Dim coord_array(1)
    point_2d_instance.GetCoordinates coord_array
    get_coordinates = coord_array
End Function
"""

system_service = CATIA.Application.SystemService


drawingDoc = CATIA.ActiveDocument
sheet = drawingDoc.Sheets.ActiveSheet
View = sheet.Views.ActiveView
texts_collection = View.Texts
fac2d = View.Factory2D
elements = View.GeometricElements
dic_point= {}
dic_line = {}
dic_point_label = {}
dic_line_label= {}
G = nx.Graph()
for i in range(1, elements.Count + 1):
    ele_ = elements.Item(i)
    obj_name = ele_.Name
    if ele_.GeometricType == catGeoTypePoint2D_:
        coord_array = system_service.evaluate(vba_code, 0, vba_function_name, [ele_])        
        dic_point[obj_name] = [coord_array[0], coord_array[1]]
    elif ele_.GeometricType == catGeoTypeLine2D_:
        start_point_name = ele_.StartPoint.Name
        end_point_name = ele_.EndPoint.Name
        dic_line[obj_name] = [start_point_name, end_point_name]
for i in range(1, texts_collection.Count + 1):
    text_ = texts_collection.Item(i)
    x = text_.x
    y = text_.y
    text_content = text_.Text
    try:
        ele_ = text_.AssociativeElement
    except Exception as e:
        print(e)
        continue
    obj_name = ele_.Name
    if ele_.GeometricType == catGeoTypePoint2D_:        
        dic_point_label[obj_name] = [text_content, x, y]
    elif ele_.GeometricType == catGeoTypeLine2D_:
        dic_line_label[obj_name] = [x, y]
        dic_line[obj_name].append(float(text_content))
dic_whole = {"dic_point": dic_point,
             "dic_point_label": dic_point_label,
             "dic_line": dic_line,
             "dic_line_label": dic_line_label}
with open("graph.pickle", 'wb') as f:
    pickle.dump(dic_whole, f)

# G.add_edge("branch_point", "tx101", weight=103)
